//
// Created by zhangcc on 2021/5/31.
//

#include "fastrt/reid_interface.h"
#include "fastrt/ioHelper.h"

static const fastrt::FastreidBackboneType BACKBONE = fastrt::FastreidBackboneType::r50;
static const fastrt::FastreidHeadType HEAD = fastrt::FastreidHeadType::EmbeddingHead;
static const fastrt::FastreidPoolingType HEAD_POOLING = fastrt::FastreidPoolingType::gempoolP;
static const int LAST_STRIDE = 1;
static const bool WITH_IBNA = true;
static const bool WITH_NL = true;
static const int EMBEDDING_DIM = 0;
static const int OUTPUTSIZE = 2048;

int FastREID::init(const std::string& weights){
    int device_count = 0;
    cudaGetDeviceCount(&device_count);
    assert(device_count > 0);
    if((device_id + 1) > device_count)
    {
        device_id = 0;
    }
    std::cout << "GPU_nums: " << device_count << std::endl;
    std::cout << "FastREID GPU ID: " << device_id << std::endl;
    cudaSetDevice(device_id);

    std::string engine = weights.substr(0, weights.find_last_of(".")) + ".engine";
    if (handle) {
        delete handle;
        handle = nullptr;
    }
    trt::ModelConfig modelCfg {
            weights, engine, max_batch_size, size.height, size.width, OUTPUTSIZE, device_id};

    fastrt::FastreidConfig reidCfg {
            BACKBONE,
            HEAD,
            HEAD_POOLING,
            LAST_STRIDE,
            WITH_IBNA,
            WITH_NL,
            EMBEDDING_DIM};

    handle = (void *)(new fastrt::Baseline{modelCfg});

    std::string buffer = nvinfer1::readBuffer(engine.c_str());
    if (!buffer.size() || alway_serialize)
    {
        fastrt::ModuleFactory moduleFactory;
        fastrt::Baseline *baseline = (fastrt::Baseline *)handle;
        std::cout << "[Serializling Engine] " << engine << std::endl;
        if (!baseline->serializeEngine(engine,
                                      {std::move(moduleFactory.createBackbone(reidCfg)),
                                       std::move(moduleFactory.createHead(reidCfg))}))
        {
            std::cout << "SerializeEngine Failed." << std::endl;
            return -1;
        }
    }

    std::cout << "[Deserializling Engine]" << std::endl;
    fastrt::Baseline *baseline = ((fastrt::Baseline *)handle);
    if (!baseline->deserializeEngine(engine))
    {
        std::cout << "DeserializeEngine Failed." << std::endl;
        return -1;
    }
    std::cout << "[DeserializeEngine Success]" << std::endl;

    return 0;
}

std::vector<std::vector<float>> FastREID::extract(std::vector<cv::Mat> &images)
{
    cudaSetDevice(device_id);

    fastrt::Baseline * baseline = (fastrt::Baseline *)handle;
    std::vector<std::vector<float> > feats;
	
	int batch_size = images.size();
    int height = images[0].rows;
    int width = images[0].cols;

//    std::cout << "batch_size: " << batch_size << std::endl;

    int iters = ceil(1.0f * batch_size / max_batch_size);
//    std::cout << "iter: " << iters << std::endl;
    std::vector<cv::Mat> images_cpy;
    for (size_t i = 0; i < iters * max_batch_size; i++)
    {
        if (i < batch_size)
        {
            images_cpy.push_back(images[i]);
        }
        else
        {
            images_cpy.emplace_back(cv::Mat(width, height, CV_8UC3, cv::Scalar(0, 0, 0)));
        }
    }

    for (size_t i = 0; i < iters; i++)
    {

        std::vector<cv::Mat>::const_iterator First = images_cpy.begin() + i * max_batch_size; // 找到第二个迭代器
        std::vector<cv::Mat>::const_iterator Second = images_cpy.begin() + (i + 1) * max_batch_size; // 找到第三个迭代器
        std::vector <cv::Mat> inputs;
        inputs.assign(First, Second);

        if (!baseline->inference(inputs))
        {
            std::cout << "Inference Failed." << std::endl;
        }
        float *feat_embedding = baseline->getOutput();
        for (size_t img_idx = 0; img_idx < max_batch_size; ++img_idx) {
            if (i*max_batch_size+img_idx > batch_size-1)
                break;

            std::vector<float> feat(OUTPUTSIZE);
            memcpy((void *)feat.data(), (feat_embedding+img_idx*OUTPUTSIZE), OUTPUTSIZE*sizeof(float));
            feats.emplace_back(feat);
        }
    }
    
    return feats;
}

FastREID::~FastREID()
{
    fastrt::Baseline * baseline = (fastrt::Baseline *)handle;
    delete baseline;
    handle = nullptr;
}